import h5py
import torch
from torch.utils.data import DataLoader
from torch.utils.data import TensorDataset
import numpy as np
import pywt

# 加入小波分解
def load_ascad(data_file_path, bs, shuffle_, mode, wave_flag=False, get_level=1, max_level=5):
    in_file = h5py.File(data_file_path, 'r')
    plain_text_ = None
    if (mode == 0):
        traces_profiling = np.array(in_file['Profiling_traces/traces'])
        labels_profiling = np.array(in_file['Profiling_traces/labels'])
        metadatas = in_file['Profiling_traces/metadata']
        plain_text = []
        for md in metadatas:
            plain_text.append(md[0][2])
        plain_text_ = np.array(plain_text)
    else:
        traces_profiling = np.array(in_file['Attack_traces/traces'])
        labels_profiling = np.array(in_file['Attack_traces/labels'])
        metadatas = in_file['Attack_traces/metadata']
        plain_text = []
        for md in metadatas:
            plain_text.append(md[0][2])
        plain_text_ = np.array(plain_text)
    plain_text_ = torch.from_numpy(plain_text_)
    if wave_flag:
        traces_profiling = np.array([wavelet_decompose(traces_profiling[i], get_level, max_level) for i in range(len(traces_profiling))])
        print(traces_profiling.shape)
    traces, labels = torch.from_numpy(traces_profiling).float(), torch.from_numpy(labels_profiling)
    ascad_dataset = TensorDataset(traces, labels)
    loader = DataLoader(dataset=ascad_dataset, batch_size=bs, shuffle=shuffle_)
    return loader, plain_text_

def wavelet_decompose(data, get_level, max_level=5):
    result = []
    coeffs = pywt.wavedec(data, 'db4', level=max_level)
    cA5, cD5, cD4, cD3, cD2, cD1 = coeffs
    # 小波重构
    result.append(pywt.waverec(np.multiply(coeffs, [1, 0, 0, 0, 0, 0]).tolist(), 'db4'))
    result.append(pywt.waverec(np.multiply(coeffs, [0, 1, 0, 0, 0, 0]).tolist(), 'db4'))
    result.append(pywt.waverec(np.multiply(coeffs, [0, 0, 1, 0, 0, 0]).tolist(), 'db4'))
    result.append(pywt.waverec(np.multiply(coeffs, [0, 0, 0, 1, 0, 0]).tolist(), 'db4'))
    result.append(pywt.waverec(np.multiply(coeffs, [0, 0, 0, 0, 1, 0]).tolist(), 'db4'))
    result.append(pywt.waverec(np.multiply(coeffs, [1, 0, 0, 0, 0, 1]).tolist(), 'db4'))
    return result[get_level]
